package ru.ifmo.swing.genetics.monitoring.common;

import org.jfree.chart.ChartFactory;
import org.jfree.chart.ChartPanel;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.axis.NumberAxis;
import org.jfree.chart.axis.ValueAxis;
import org.jfree.chart.plot.PlotOrientation;
import org.jfree.data.xy.XYSeries;
import org.jfree.data.xy.XYSeriesCollection;
import org.uncommons.watchmaker.framework.PopulationData;
import org.uncommons.watchmaker.framework.islands.IslandEvolutionObserver;

import javax.swing.*;
import java.awt.*;
import java.awt.event.ItemEvent;
import java.awt.event.ItemListener;
import java.util.ArrayList;
import java.util.Collection;

public class ChartObserver<T> implements IslandEvolutionObserver<T>, SwingEvolutionObserver<T> {
    private static final int SHOW_FIXED_GENERATIONS = 200;
    private final JPanel view;

    private final java.util.List<SinglePlot<T>> plots = new ArrayList<SinglePlot<T>>();

    private final XYSeriesCollection dataSet = new XYSeriesCollection();
    private final ValueAxis domainAxis;
    private final ValueAxis rangeAxis;

    private final JRadioButton allDataButton = new JRadioButton("All Data", false);
    private final JFreeChart chart;

    private double maxY = 1;
    private double minY = 0;


    ChartObserver(boolean islands, String title, Collection<ChartDataProvider<T>> dataProviders) {
        view = new JPanel(new BorderLayout());
        for (ChartDataProvider<T> dataProvider : dataProviders) {
            SinglePlot<T> plot = new SinglePlot<T>(dataProvider);
            plots.add(plot);
            dataSet.addSeries(plot);
        }
        chart = ChartFactory.createXYLineChart(title,
                islands ? "Epochs" : "Generations",
                "",
                dataSet,
                PlotOrientation.VERTICAL,
                true, // Legend.
                false, // Tooltips.
                false);
        this.domainAxis = chart.getXYPlot().getDomainAxis();
        this.rangeAxis = chart.getXYPlot().getRangeAxis();
        domainAxis.setStandardTickUnits(NumberAxis.createIntegerTickUnits());
        domainAxis.setLowerMargin(0);
        domainAxis.setUpperMargin(0.05);
        domainAxis.setRangeWithMargins(0, SHOW_FIXED_GENERATIONS);
        rangeAxis.setRange(minY, maxY);
        ChartPanel chartPanel = new ChartPanel(chart,
                ChartPanel.DEFAULT_WIDTH,
                ChartPanel.DEFAULT_HEIGHT,
                ChartPanel.DEFAULT_MINIMUM_DRAW_WIDTH,
                ChartPanel.DEFAULT_MINIMUM_DRAW_HEIGHT,
                ChartPanel.DEFAULT_MAXIMUM_DRAW_WIDTH,
                ChartPanel.DEFAULT_MAXIMUM_DRAW_HEIGHT,
                false, // Buffered
                false, // Properties
                true, // Save
                true, // Print
                false, // Zoom
                false); // Tooltips
        view.add(chartPanel, BorderLayout.CENTER);
        view.add(createControls(islands), BorderLayout.SOUTH);
    }


    /**
     * Creates the GUI controls for toggling graph display options.
     *
     * @return A component that can be added to the main panel.
     */
    private JComponent createControls(boolean islands) {
        JPanel controls = new JPanel(new FlowLayout(FlowLayout.RIGHT));

        allDataButton.addItemListener(new ItemListener() {
            public void itemStateChanged(ItemEvent ev) {
                updateDomainAxisRange();
            }
        });
        String text = "Last " + SHOW_FIXED_GENERATIONS + (islands ? " Epochs" : " Generations");
        JRadioButton recentDataButton = new JRadioButton(text, true);
        ButtonGroup buttonGroup = new ButtonGroup();
        buttonGroup.add(allDataButton);
        buttonGroup.add(recentDataButton);

        controls.add(allDataButton);
        controls.add(recentDataButton);

        return controls;
    }


    /**
     * If "all data" is selected, set the range of the domain axis to include all
     * values.  Otherwise set it to show the most recent 200 generations.
     */
    private void updateDomainAxisRange() {
        int count = dataSet.getSeries(0).getItemCount();
        if (count < SHOW_FIXED_GENERATIONS) {
            domainAxis.setRangeWithMargins(0, SHOW_FIXED_GENERATIONS);
        } else if (allDataButton.isSelected()) {
            domainAxis.setRangeWithMargins(0, Math.max(SHOW_FIXED_GENERATIONS, count));
        } else {
            domainAxis.setRangeWithMargins(count - SHOW_FIXED_GENERATIONS, count);
        }
    }


    /**
     * {@inheritDoc}
     */
    public void populationUpdate(final PopulationData<? extends T> populationData) {
        SwingUtilities.invokeLater(new Runnable() {
            public void run() {
                chart.setNotify(false); // Avoid triggering a redraw for every change we make in this method.
                if (populationData.getGenerationNumber() == 0) {
                    // The graph might be showing data from a previous run, so clear it.
                    for (SinglePlot<T> plot : plots) {
                        plot.clear();
                    }
                }
                double high = maxY;
                double low = minY;
                for (SinglePlot<T> plot : plots) {
                    double value = plot.populationUpdate(populationData);
                    high = Math.max(high, value);
                    low = Math.min(low, value);
                }

                if (high > maxY || low < minY) {
                    maxY = high;
                    minY = low;
                }
                rangeAxis.setRange(minY - Math.abs(minY * 0.05), maxY + Math.abs(maxY * 0.05));

                updateDomainAxisRange();
                chart.setNotify(true); // Redraw all at once now.
            }
        });
    }


    public void islandPopulationUpdate(int islandIndex, PopulationData<? extends T> populationData) {
        // Do nothing.
    }

    public JComponent getView() {
        return view;
    }

    private class SinglePlot<T> extends XYSeries {
        private final ChartDataProvider<T> dataProvider;

        public SinglePlot(ChartDataProvider<T> dataProvider) {
            super(dataProvider.getName());
            this.dataProvider = dataProvider;
        }

        public double populationUpdate(PopulationData<? extends T> populationData) {
            double value = dataProvider.getValue(populationData);
            add(populationData.getGenerationNumber(), value);
            return value;
        }
    }

}
